
from langchain.retrievers import ParentDocumentRetriever
from langchain.schema import Document
from langchain_community.document_transformers import LongContextReorder
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

from Knowledge.milvus_store import vectorstore
from Knowledge.pg_store import PGStore

# 向量模型：使用 bge-m3，输出维度为 1024
knowledge_embedder = OllamaEmbeddings(model="bge-m3:latest", base_url="http://192.168.7.3:11434")
qa_model = OllamaEmbeddings(model="nomic-embed-text:latest", base_url="http://192.168.7.3:11434")

# 父文档存储 + 子文档向量索引
pg_store = PGStore()

retriever = ParentDocumentRetriever(
    vectorstore=vectorstore,
    docstore=pg_store,
    child_splitter=RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100),
)


def knowledge_search(query, top_k=2):
    """
    1. 查询向量化
    2. 使用 Milvus 检索相似子文档
    3. 用 parent_id 获取完整父文档
    4. 使用 LongContextReorder 对父文档重排
    """
    # ✅ Step 1: 向量化查询
    query_embedding = knowledge_embedder.embed_query("为信息检索优化的问句: " + query)

    # ✅ Step 2: 在 Milvus 检索子文档
    search_results = vectorstore.similarity_search_by_vector(
        embedding=query_embedding, k=top_k
    )

    if not search_results:
        print("⚠️ 没有找到相关文档")
        return []

    # ✅ Step 3: 提取 parent_id 并查找完整文档
    parent_ids = list(set(doc.metadata["parent_id"] for doc in search_results))
    parent_docs_dict = retriever.docstore.mget(parent_ids)

    # ✅ Step 4: 转为 Document 类型，并排序
    doc_list = [
        Document(page_content=content, metadata={"parent_id": pid})
        for pid, content in parent_docs_dict.items()
    ]

    reorderer = LongContextReorder()
    reordered_docs = reorderer.transform_documents(doc_list)

    return [doc.page_content for doc in reordered_docs]


def qa_search(query, top_k=2):
    query_embedding = qa_model.embed_query("为信息检索优化的问句: " + query)
    search_results = vectorstore.similarity_search_by_vector(
        embedding=query_embedding, k=top_k
    )
    if not search_results:
        print("⚠️ 没有找到相关文档")
        return []
    return search_results
